{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZloPIuRHn97X"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors. [Licensed under the Apache License, Version 2.0](#scrollTo=Afd8bu4xJOgh)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tNgCmfUvJNoF"
      },
      "outputs": [],
      "source": [
        "// #@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "// you may not use this file except in compliance with the License.\n",
        "// You may obtain a copy of the License at\n",
        "//\n",
        "// https://www.apache.org/licenses/LICENSE-2.0\n",
        "//\n",
        "// Unless required by applicable law or agreed to in writing, software\n",
        "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "// See the License for the specific language governing permissions and\n",
        "// limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AlvdCHw5JGyx"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/swift/tutorials/custom_differentiation\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">TensorFlow.org で表示</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/swift/tutorials/custom_differentiation.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab で実行</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/swift/tutorials/custom_differentiation.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub でソースを表示</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c_1u7JSBMx3x"
      },
      "source": [
        "# カスタム微分\n",
        "\n",
        "このチュートリアルでは、わずか 5 行の Swift によって独自のカスタム導関数を定義し、導関数の手直しを行い、独自の勾配チェックポイント API を実装する方法を紹介します。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gHuQo_kCTjFr"
      },
      "source": [
        "## カスタム導関数を宣言する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LP0gMw56TlvH"
      },
      "source": [
        "微分可能なパラメータと結果を持つ、任意の Swift 関数のカスタム導関数を定義することができます。これにより、C 関数をインポートして微分可能にすることもできます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j0a8prgZTlEO"
      },
      "outputs": [],
      "source": [
        "import Glibc\n",
        "\n",
        "func sillyExp(_ x: Float) -> Float {\n",
        "    let 𝑒 = Float(M_E)\n",
        "    print(\"Taking 𝑒(\\(𝑒)) to the power of \\(x)!\")\n",
        "    return pow(𝑒, x)\n",
        "}\n",
        "\n",
        "@derivative(of: sillyExp)\n",
        "func sillyDerivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {\n",
        "    let y = sillyExp(x)\n",
        "    return (value: y, pullback: { v in v * y })\n",
        "}\n",
        "\n",
        "print(\"exp(3) =\", sillyExp(3))\n",
        "print(\"𝛁exp(3) =\", gradient(of: sillyExp)(3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eQPX9r3R5OP-"
      },
      "source": [
        "## 導関数の伝播を防ぐ\n",
        "\n",
        "機械学習のユースケースで「勾配を止める」として一般的に知られている `withoutDerivative(at:)` メソッドは、導関数の伝播を防ぎます。\n",
        "\n",
        "さらに、`withoutDerivative(at:)` は微分してはいけないものを識別し、より効率的な導関数を生成するため、 Swift コンパイラの役に立つ場合があります。関数の微分が常にゼロになることを検出可能な場合は Swift コンパイラが警告を生成します。明示的に `withoutDerivative(at:)` を使用すると、その警告は非表示になります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ctRt6vBO5Wle"
      },
      "outputs": [],
      "source": [
        "let x: Float = 2.0\n",
        "let y: Float = 3.0\n",
        "gradient(at: x, y) { x, y in\n",
        "    sin(sin(sin(x))) + withoutDerivative(at: cos(cos(cos(y))))\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EeV3wXQ79WS2"
      },
      "source": [
        "## 導関数の手直し\n",
        "\n",
        "[`withDerivative(_:)`](https://www.tensorflow.org/swift/api_docs/Protocols/Differentiable#/s:10TensorFlow14DifferentiablePAAE12withGradientyxy15CotangentVectorQzzcF) メソッドは、囲む関数のバックプロパゲーション中の値で、任意の演算（突然変異を含む）を勾配上で実行させます。\n",
        "\n",
        "これを使用して、デバッグを行ったり実験的にバックプロパゲーションを調整したりすることができます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AHV0ryTiD6j8"
      },
      "source": [
        "### どこでも動作する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9zKSeUjTmbxq"
      },
      "source": [
        "標準ライブラリが提供しているすべての微分 API は、 `Differentiable` プロトコルに準拠するすべての型を一般的に定義します。これには `Float`、`Double`、`Float80`、SIMD ベクトル、そして独自の型も含まれます！\n",
        "\n",
        "`Differentiable` プロトコルの詳細については、テクニカルドキュメントの[可微分の型](https://github.com/tensorflow/swift/blob/master/docs/DifferentiableTypes.md)をご覧ください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eKne7szjD8lr"
      },
      "outputs": [],
      "source": [
        "var x: Float = 30\n",
        "gradient(at: x) { x -> Float in\n",
        "    // Print the partial derivative with respect to the result of `sin(x)`.\n",
        "    let a = sin(x).withDerivative { print(\"∂+/∂sin = \\($0)\") } \n",
        "    // Force the partial derivative with respect to `x` to be `0.5`.\n",
        "    let b = log(x.withDerivative { (dx: inout Float) in\n",
        "        print(\"∂log/∂x = \\(dx), but rewritten to 0.5\");\n",
        "        dx = 0.5\n",
        "    })\n",
        "    return a + b\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vmw0gkqlD9xf"
      },
      "source": [
        "### ニューラルネットワークモジュール内で使用する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JCf_OplsWzhW"
      },
      "source": [
        "単純な `Float` 関数で使うのと同じように、これはあらゆる数値的アプリケーションで使用することができます。[Swift for TensorFlow ディープラーニングライブラリ](https://github.com/tensorflow/swift-apis)を使用して構築された以下のニューラルネットワークもその 1 つです。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fnSeAbs9-hf3"
      },
      "outputs": [],
      "source": [
        "import TensorFlow\n",
        "\n",
        "struct MLP: Layer {\n",
        "    var layer1 = Dense<Float>(inputSize: 2, outputSize: 10, activation: relu)\n",
        "    var layer2 = Dense<Float>(inputSize: 10, outputSize: 1, activation: relu)\n",
        "    \n",
        "    @differentiable\n",
        "    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {\n",
        "        let h0 = layer1(input).withDerivative { print(\"∂L/∂layer1 =\", $0) }\n",
        "        return layer2(h0)\n",
        "    }\n",
        "}\n",
        "\n",
        "var classifier = MLP()\n",
        "let optimizer = SGD(for: classifier, learningRate: 0.02)\n",
        "\n",
        "let x: Tensor<Float> = [[0, 0], [0, 1], [1, 0], [1, 1]]\n",
        "let y: Tensor<Float> = [0, 1, 1, 0]\n",
        "\n",
        "for _ in 0..<10 {\n",
        "    let 𝛁model = gradient(at: classifier) { classifier -> Tensor<Float> in\n",
        "        let ŷ = classifier(x).withDerivative { print(\"∂L/∂ŷ =\", $0) }\n",
        "        let loss = (ŷ - y).squared().mean()\n",
        "        print(\"Loss: \\(loss)\")\n",
        "        return loss\n",
        "    }\n",
        "    optimizer.update(&classifier, along: 𝛁model)\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TzLfTj28gEUD"
      },
      "source": [
        "## バックプロパゲーション中の活性化を再計算してメモリを節約する（チェックポインティング）\n",
        "\n",
        "チェックポインティングは、メモリを節約するためにリバースモード自動微分をする従来の手法です。元の計算で大きな中間値を保存して導関数計算をするのではなく、必要に応じてバックプロパゲーション中に中間値を再計算します。\n",
        "\n",
        "この手法は、最近のディープラーニングライブラリでも実現されています。Swift では、[`withRecomputationInPullbacks(_:)`](https://www.tensorflow.org/swift/api_docs/Protocols/Differentiable#/s:10TensorFlow14DifferentiablePAAE28withRecomputationInPullbacksyqd__qd__xcAaBRd__lF) API によってバックプロパゲーション中に何を再計算するかを制御することができ、すべての `Differentiable` 型で利用可能です。\n",
        "\n",
        "しかしここでは、数行のコードのみを使用して独自の勾配チェックポインティング API を新規に定義する方法を見ていきます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5cZe-JbjwMfZ"
      },
      "source": [
        "### 勾配チェックポインティング API"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "606ob1dn2v77"
      },
      "source": [
        "独自の勾配チェックポインティング API である `makeRecomputedInGradient(_:)` を定義することができますが、標準ライブラリ関数 [`differentiableFunction(from:)`](https://www.tensorflow.org/swift/api_docs/Functions#/s:10TensorFlow22differentiableFunction4fromq0_x_q_tcq0_5value_15CotangentVectorQz_AEQy_tAEQy0_c8pullbacktx_q_tc_tAA14DifferentiableRzAaJR_AaJR0_r1_lF) の観点から見ると、これは微分関数（「ベクトルヤコビアン積（VJP）関数」とも呼ばれる）から微分可能な関数を直接作成するための省略表現です。\n",
        "\n",
        "これまで説明したように、微分関数は元の関数の結果とプルバッククロージャのタプルを返します。`original(x)` を `value:` で返し、`original` 上で`pullback(at:in:)` を呼び出すことにより、元の関数を再評価してプルバックを取得します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b1uU3tcVwl_1"
      },
      "outputs": [],
      "source": [
        "/// Given a differentiable function, returns the same differentiable function except when\n",
        "/// derivatives of this function are being computed. In that case, values in the original function needed\n",
        "/// for computing the derivatives will be recomputed, instead of being captured by the differential or pullback.\n",
        "///\n",
        "/// - Parameter body: The body of the differentiable function.\n",
        "/// - Returns: The same differentiable function whose derivatives, when computed, will recompute\n",
        "///   some values from the original function.\n",
        "func makeRecomputedInGradient<T: Differentiable, U: Differentiable>(\n",
        "    _ original: @escaping @differentiable (T) -> U\n",
        ") -> @differentiable (T) -> U {\n",
        "    return differentiableFunction { x in\n",
        "        (value: original(x), pullback: { v in pullback(at: x, in: original)(v) })\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UbeKj7NEF7zz"
      },
      "source": [
        "### 動作を確認する"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oee8SXital45"
      },
      "outputs": [],
      "source": [
        "let input: Float = 10.0\n",
        "print(\"Running original computation...\")\n",
        "\n",
        "// Differentiable multiplication with checkpointing.\n",
        "let square = makeRecomputedInGradient { (x: Float) -> Float in\n",
        "    print(\"  Computing square...\")\n",
        "    return x * x\n",
        "}\n",
        "\n",
        "// Differentiate `f(x) = (cos(x))^2`.\n",
        "let (output, backprop) = valueWithPullback(at: input) { input -> Float in\n",
        "    return square(cos(input))\n",
        "}\n",
        "print(\"Running backpropagation...\")\n",
        "let grad = backprop(1)\n",
        "print(\"Gradient = \\(grad)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7SxWsSUqF9Bh"
      },
      "source": [
        "### ニューラルネットワークモジュールに拡張する\n",
        "\n",
        "この例では、単純な畳み込みニューラルネットワークを定義します。\n",
        "\n",
        "```swift\n",
        "struct Model: Layer {     var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6))     var maxPool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))     var flatten = Flatten<Float>()     var dense = Dense<Float>(inputSize: 36 * 6, outputSize: 10)      @differentiable     func call(_ input: Tensor<Float>) -> Tensor<Float> {         return input.sequenced(through: conv, maxPool, flatten, dense)     } }\n",
        "```\n",
        "\n",
        "バックプロパゲーション中に畳み込みレイヤー（`conv`）の活性化を再計算させる必要があります。しかし、`makeRecomputedInGradient(_:)` を使用すると、特に [`sequenced(in:through:_:_:_:_:)`](https://www.tensorflow.org/swift/api_docs/Protocols/Differentiable#/s:10TensorFlow14DifferentiablePAAE9sequenced2in7through____6OutputQyd_3_AA7ContextC_qd__qd_0_qd_1_qd_2_qd_3_t5InputQyd__RszAA5LayerRd__AaMRd_0_AaMRd_1_AaMRd_2_AaMRd_3_AKQyd_0_AGRtd__AKQyd_1_AGRtd_0_AKQyd_2_AGRtd_1_AKQyd_3_AGRtd_2_r3_lF) を用いて順次レイヤーを適用する場合は、結果のコードが扱いにくく見える可能性があります。\n",
        "\n",
        "```swift\n",
        "input.sequenced(in: context, through: conv, maxPool, flatten, dense)\n",
        "```\n",
        "\n",
        "そこで、バックプロパゲーション中にレイヤーをラップして活性化を再計算させる、**特殊レイヤー型**を定義してみてはどうでしょう。早速、これを試してみます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZP86M5RjP3OG"
      },
      "source": [
        "まず、二値関数を取る `makeRecomputedInGradient(_:)` 関数を定義します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bEm-n5H0QB8s"
      },
      "outputs": [],
      "source": [
        "// Same as the previous `makeRecomputedInGradient(_:)`, except it's for binary functions.\n",
        "func makeRecomputedInGradient<T: Differentiable, U: Differentiable, V: Differentiable>(\n",
        "    _ original: @escaping @differentiable (T, U) -> V\n",
        ") -> @differentiable (T, U) -> V {\n",
        "    return differentiableFunction { x, y in\n",
        "        (value: original(x, y), pullback: { v in pullback(at: x, y, in: original)(v) })\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YU6DgqXxP5Nl"
      },
      "source": [
        "次に、汎用レイヤー `ActivationDiscarding<Wrapped>` を定義します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ao1r_lIPGeOl"
      },
      "outputs": [],
      "source": [
        "import TensorFlow\n",
        "\n",
        "/// A layer wrapper that makes the underlying layer's activations be discarded during application\n",
        "/// and recomputed during backpropagation.\n",
        "struct ActivationDiscarding<Wrapped: Layer>: Layer {\n",
        "    /// The wrapped layer.\n",
        "    var wrapped: Wrapped\n",
        "\n",
        "    @differentiable\n",
        "    func callAsFunction(_ input: Wrapped.Input) -> Wrapped.Output {\n",
        "        let apply = makeRecomputedInGradient { (layer: Wrapped, input: Input) -> Wrapped.Output in\n",
        "            print(\"    Applying \\(Wrapped.self) layer...\")\n",
        "            return layer(input)\n",
        "        }\n",
        "        return apply(wrapped, input)\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HqPXwwuTRjmz"
      },
      "source": [
        "最後に、アプリケーション中に活性化が破棄されバックプロパゲーション中に再計算されることを除き、同じレイヤーを返すメソッドを全てのレイヤーに追加することができます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PGgkNnNNR1th"
      },
      "outputs": [],
      "source": [
        "extension Layer {\n",
        "    func discardingActivations() -> ActivationDiscarding<Self> {\n",
        "        return ActivationDiscarding(wrapped: self)\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8PP-NZ9XU5_n"
      },
      "source": [
        "モデルに関しては、ここで変更が必要なのは畳み込みレイヤーを活性化再計算レイヤーにラップすることのみです。\n",
        "\n",
        "```swift\n",
        "var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6)).discardingActivations()\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bCwNPtCfSbGi"
      },
      "source": [
        "あとは、それをモデルで使用するだけです！"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gsWGwFjOJ3Md"
      },
      "outputs": [],
      "source": [
        "struct Model: Layer {\n",
        "    var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6)).discardingActivations()\n",
        "    var maxPool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))\n",
        "    var flatten = Flatten<Float>()\n",
        "    var dense = Dense<Float>(inputSize: 36 * 6, outputSize: 10)\n",
        "\n",
        "    @differentiable\n",
        "    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {\n",
        "        return input.sequenced(through: conv, maxPool, flatten, dense)\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmFxciU6VYdF"
      },
      "source": [
        "トレーニングループを実行すると、畳み込みレイヤーの活性化が 2 回（レイヤー適用中に 1 回、バックプロパゲーション中に 1 回）計算されていることが分かります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-x1nYu0uVSPn"
      },
      "outputs": [],
      "source": [
        "// Use random training data.\n",
        "let x = Tensor<Float>(randomNormal: [10, 16, 16, 3])\n",
        "let y = Tensor<Int32>(rangeFrom: 0, to: 10, stride: 1)\n",
        "\n",
        "var model = Model()\n",
        "let opt = SGD(for: model)\n",
        "\n",
        "for i in 1...5 {\n",
        "    print(\"Starting training step \\(i)\")\n",
        "    print(\"  Running original computation...\")\n",
        "    let (logits, backprop) = model.appliedForBackpropagation(to: x)\n",
        "    let (loss, dL_dŷ) = valueWithGradient(at: logits) { logits in\n",
        "        softmaxCrossEntropy(logits: logits, labels: y)\n",
        "    }\n",
        "    print(\"  Loss: \\(loss)\")\n",
        "    print(\"  Running backpropagation...\")\n",
        "    let (dL_dθ, _) = backprop(dL_dŷ)\n",
        "    \n",
        "    opt.update(&model, along: dL_dθ)\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gzRaZLa_WX0u"
      },
      "source": [
        "このように、一般的な可微分プログラミングライブラリを簡単に異なるドメインに定義できます。"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "custom_differentiation.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Swift",
      "name": "swift"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
